from sage.all import *
import random
from ast import literal_eval

from sage.crypto.mq import sr
from sage.interfaces.magma import magma
from sage.rings.polynomial.pbori import BooleanPolynomialIdeal as BPIdeal
from sage.rings.polynomial.polynomial_ring_constructor import \
    BooleanPolynomialRing_constructor as BPRing
from sage.rings.polynomial.pbori import substitute_variables

from colorama import Back, Style


class PS():
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)

    def _variety(self, unique_polys_only=False):
        text_output = False
        orders = {'lex': 'lex', 'deglex': 'glex', 'degrevlex': 'grevlex'}
        # TODO: Use [text := text.replace(a, b) for a, b in replacements][-1]
        # instead of plain .replace to avoid repetition
        BPR_m = ('P' +
                 str(self.ring.gens()).replace('(', '<').replace(')', '>') +
                 ' := BooleanPolynomialRing(' + str(self.ring.n_variables()) +
                 ', "' + orders[self.order] + '");\n\n')
        magma.eval(BPR_m)
        polys_m = []
        polys_txt = ''
        polys = [p for s in self.polys for p in s]
        if unique_polys_only:
            polys = set(polys)
        for poly, i in zip(polys, range(len(polys))):
            i = str(i)
            monomials = poly.monomials()
            p_a = ' + '.join(str(m) for m in monomials[:len(monomials) // 2])
            p_b = ' + '.join(str(m) for m in monomials[len(monomials) // 2:])
            magma.eval('p_a' + i + ' := ' + p_a + ';\n')
            magma.eval('p_b' + i + ' := ' + p_b + ';\n')
            magma.eval('p' + i + ' := p_a' + i + ' + p_b' + i + ';\n')
            polys_m += ['p' + i]
            if text_output:
                polys_txt += 'p' + i + ' := ' + str(poly) + ';\n\n'
        print('Attempting to compute the gb.')
        I_m = 'I := Ideal(' + str(polys_m) + ');\n\n GroebnerBasis(I);\n'
        if text_output:
            with open('magma.txt', 'w') as f:
                f.write(BPR_m + polys_txt + I_m)
        self.gb = magma.eval(I_m)
        print('Attempting to obtain the variety.')
        self.V = literal_eval(
            magma.eval('Variety(I)').replace('<', '(').replace('>', ')'))

    def keys(self, unique_polys_only=False):
        try:
            self.V
        except AttributeError:
            self._variety(unique_polys_only)
        finally:
            self._keys = [v[-len(list(self.key_orig)):] for v in self.V]
            return self._keys


class SR(sr.SR_gf2):
    def __init__(self, n=1, r=1, c=1, e=4, star=False, **kwargs):
        super().__init__(n, r, c, e, star, gf2=True, **kwargs)
        self.order = kwargs.get('order', 'lex')
        self.s_size = r * c * e

    def _ps_raw(self, n=1):
        ps_raw = {}
        for _ in range(2**7):
            ps_raw['pt'] = []
            ps_raw['ct'] = []
            ps_raw['ps'] = []
            ps_raw['key_orig'] = self.vector(
                random.choices(range(2), k=self.s_size))
            for _ in range(n):
                for _ in range(2**7):
                    pt = self.vector(random.choices(range(2), k=self.s_size))
                    if pt in ps_raw['pt']:
                        continue
                    try:
                        ct = self(P=pt, K=ps_raw['key_orig'])
                        ps, _ = self.polynomial_system(P=pt, C=ct)
                        ps_raw['pt'] += [pt]
                        ps_raw['ct'] += [ct]
                        ps_raw['ps'] += [ps]
                        break
                    except ZeroDivisionError:
                        continue
                else:
                    break
            else:
                break
        else:
            raise ValueError(
                'The polynomial system could not be created either due to too '
                'many 0-inversions, or you are requesting too many polynomial '
                'systems.')
        ps_raw['ring'] = ps.ring()
        ps_raw['order'] = self.order
        return ps_raw

    def ps_key_vars(self, n=1):
        def sboxes_gb(sboxes):
            return (BPIdeal(
                BPRing(names=sbox[0].variables(),
                       order='deglex(4), deglex(4)'),
                sbox).groebner_basis()[:4]
                    for sbox in (sboxes[i:i + 12]
                                 for i in range(0, len(sboxes), 12)))

        def separate_lm(polys):
            return dict((str((p.lm())), p.subs({p.lm(): 0})) for p in polys)

        def subs_vars(polys, vars_, ring):
            if not polys:
                return []
            else:
                gens = [
                    ring(vars_[str(i)]) if str(i) in vars_ else ring(i)
                    for i in polys[0].ring().gens()
                ]
                return [substitute_variables(ring, gens, p) for p in polys]

        ps_raw = self._ps_raw(n)
        w_vars = {}
        s_vars = {}
        eliminated_vars = []
        ring = ps_raw['ring']
        ps_raw_it = iter(ps_raw['ps'].parts())
        for r_polys, ks_polys in zip(ps_raw_it, ps_raw_it):
            vars_to_subs = {}
            x_vars = {}

            # Express the key schedule in the variables of the initial key.
            for sbox in sboxes_gb(ks_polys[self.s_size:]):
                x_vars.update(separate_lm(sbox))
            eliminated_vars += list(x_vars.keys())
            s_vars = separate_lm(
                subs_vars(subs_vars(ks_polys[:self.s_size], x_vars, ring),
                          s_vars, ring))
            vars_to_subs.update(s_vars)

            # Express the round in the variables of the initial key.
            for sbox in sboxes_gb(r_polys[self.s_size:]):
                vars_to_subs.update(separate_lm(subs_vars(sbox, w_vars, ring)))
            eliminated_vars += list(w_vars.keys()) + list(vars_to_subs.keys())
            polys = subs_vars(r_polys[:self.s_size], vars_to_subs, ring)
            w_vars = separate_lm(polys)
        ps_raw['ring'] = ring.remove_var(*eliminated_vars)
        ps_raw['polys'] = polys
        return PS(**ps_raw)

    def ps_aux_vars(self, n=1):
        ps_raw = self._ps_raw(n)
        ps_raw['polys'] = [[p for s in ss.parts() for p in s]
                           for ss in ps_raw['ps']]
        return PS(**ps_raw)


r = 2
c = 2
n = 1
e = 4
order = 'lex'

rijndael = SR(n, r, c, e, order=order)
ps = rijndael.ps_aux_vars(n=2)
valid_keys = []
with sr.AllowZeroInversionsContext(rijndael):
    for key in ps.keys(unique_polys_only=False):
        for pt, ct in zip(ps.pt, ps.ct):
            if ct != rijndael(P=pt, K=rijndael.vector(key)):
                break
        else:
            valid_keys += [key]
            print(Back.GREEN + 'A valid key found.' + Style.RESET_ALL)

if not valid_keys:
    print(Back.RED + 'No valid key found.' + Style.RESET_ALL)
